import torch
from torch.utils.data.sampler import Sampler
import numpy as np


class InfiniteSampler(Sampler):
    def __init__(self, size, shuffle=True):
        self.size = size
        self.shuffle = shuffle

    def __iter__(self):
        while True:
            if self.shuffle:
                yield from torch.randperm(self.size)
            else:
                yield from torch.arange(self.size)


class InfiniteInstanceBalanceSampler(Sampler):
    def __init__(self, dataset, num_classes):
        self.weights = self.calculate_weights(dataset, num_classes)
        self.size = len(self.weights)

    def calculate_weights(self, dataset, num_classes):
        weights = np.zeros((len(dataset), num_classes), dtype=np.float32)
        for i, data in enumerate(dataset):
            for obj in data['objects']:
                weights[i, obj['category']] += 1
        weights = np.linalg.pinv(weights).sum(axis=0)
        weights_sum = weights.sum()
        assert weights_sum > 0, "unable to balance instances"
        return weights / weights_sum

    def __iter__(self):
        while True:
            yield np.random.choice(self.size, 1, p=self.weights).item()


class InfiniteCategoryBalanceSampler(Sampler):
    def __init__(self, dataset, num_classes):
        self.weights = self.calculate_weights(dataset, num_classes)
        self.size = len(self.weights)

    def calculate_weights(self, dataset, num_classes):
        weights = np.zeros((len(dataset), num_classes), dtype=np.float32)
        for i, data in enumerate(dataset):
            for obj in data['objects']:
                weights[i, obj['category']] = 1
        weights = np.linalg.pinv(weights).sum(axis=0)
        weights_sum = weights.sum()
        assert weights_sum > 0, "unable to balance categories"
        return weights / weights_sum

    def __iter__(self):
        while True:
            yield np.random.choice(self.size, 1, p=self.weights).item()
